import networkx as nx
import numpy as np
import matplotlib.pyplot as plt


def web_page(n, m, p):
    if m < 1 or n < m:
        raise nx.NetworkXError(
            f"NetworkXError must have m>1 and m<n, m={m},n={n}")

    if p > 1 or p < 0:
        raise nx.NetworkXError(f"NetworkXError p must be in [0,1], p={p}")

    G = nx.DiGraph()
    G.add_node(0)
    source = 1
    while source < n:
        G.add_node(source)
        possible_targets = np.random.choice(G.nodes(),
                                            size=min(m, G.number_of_nodes()),
                                            replace=False)

        for target in possible_targets:  # add m new links
            if np.random.random() < p:
                G.add_edge(source, target)
            else:
                neighborhood = [
                    nbr for nbr in G.neighbors(target)
                    if not G.has_edge(source, nbr) and not nbr == source
                ]
                if neighborhood:  # if there is a neighbor without a link
                    nbr = np.random.choice(neighborhood)
                    G.add_edge(source, nbr)

        source += 1
    return G


def web_page_with_random_edge(n, m, k, p):
    if m < 1 or n < m:
        raise nx.NetworkXError(
            f"NetworkXError must have m>1 and m<n, m={m},n={n}")

    if p > 1 or p < 0:
        raise nx.NetworkXError(f"NetworkXError p must be in [0,1], p={p}")

    G = nx.DiGraph()
    G.add_node(0)
    source = 1
    while source < n:
        G.add_node(source)
        if m == k:
            num_new_edge = k
        else:
            num_new_edge = np.random.randint(k, m + 1)
        possible_targets = np.random.choice(G.nodes(),
                                            size=min(num_new_edge,
                                                     G.number_of_nodes()),
                                            replace=False)

        for target in possible_targets:  # add m new links
            if np.random.random() < p:
                G.add_edge(source, target)
            else:
                neighborhood = [
                    nbr for nbr in G.neighbors(target)
                    if not G.has_edge(source, nbr) and not nbr == source
                ]
                if neighborhood:  # if there is a neighbor without a link
                    nbr = np.random.choice(neighborhood)
                    G.add_edge(source, nbr)

        source += 1
    return G


def relink(G: nx.DiGraph, p: float, m=1):
    """
    for each node link to a random other node with probability p
    """
    for source in G.nodes():
        t = list(G.nodes())
        t.remove(source)
        other_nodes = np.random.choice(t, size=min(m, len(t)), replace=False)
        for target in other_nodes:
            if np.random.random() < p and not G.has_edge(source, target):
                G.add_edge(source, target)


def web_structure(G: nx.DiGraph):
    # scc_node = nx.strongly_connected_components(G)\
    # 弱连通分量 节点集合
    wcc = list(nx.weakly_connected_components(G))
    # sorted(wcc, key=len)
    # # 最大弱连通分量
    # lwcc = wcc[-1]
    # 最大强连通分量
    lscc = max(nx.strongly_connected_components(G), key=len)
    # 包含最大强连通分量的弱连通分量  用以减小计算量 IN SCC OUT 卷须和管道 都包含于此
    # wcc_of_lscc = list(filter(lambda x:not lscc.isdisjoint(x), wcc))
    wcc_of_lscc = []
    disconnect_components = []
    for s in wcc:
        if lscc.isdisjoint(s):
            disconnect_components.append(s)
        else:
            wcc_of_lscc.append(s)

    assert len(wcc_of_lscc) == 1, "error"
    wcc_of_lscc:nx.DiGraph = G.subgraph(wcc_of_lscc[0])

    source = next(iter(lscc))
    successors = set()
    for p, c in nx.bfs_successors(wcc_of_lscc, source=source):
        successors.update(c)

    predecessors = set()
    # t = nx.DiGraph([(v, u) for u, v in wcc_of_lscc.edges()])
    t = wcc_of_lscc.reverse()
    for p, c in nx.bfs_successors(t, source=source):
        predecessors.update(c)

    assert lscc == (predecessors & successors) | {source}, "not equal"

    IN = predecessors - lscc
    OUT = successors - lscc
    Tendril = wcc_of_lscc.nodes() - (predecessors | successors)
    return lscc, IN, OUT, Tendril, disconnect_components


def plot_web_structure(G: nx.DiGraph):
    SCC, IN, OUT, Tendril, disconnect_components = web_structure(g)
    disconnect_components = set([n for i in disconnect_components for n in i])
    info = f"SCC: {len(SCC)}, IN: {len(IN)}, OUT: {len(OUT)}, Tendril: {len(Tendril)}, Disconnected_components: {len(disconnect_components)}"
    print(info)
    size = np.array([
        len(SCC),
        len(IN),
        len(OUT),
        len(Tendril),
        len(disconnect_components)
    ]) / g.number_of_nodes()
    pos0 = nx.spring_layout(g.subgraph(SCC))
    pos1 = nx.spring_layout(g.subgraph(IN))
    pos2 = nx.spring_layout(g.subgraph(OUT))
    pos3 = nx.spring_layout(g.subgraph(Tendril))
    pos4 = nx.spring_layout(g.subgraph(
        disconnect_components)) if len(disconnect_components) > 0 else {}
    for k in pos0:
        pos0[k] *= size[0]
    for k in pos1:
        pos1[k] *= size[1]
        pos1[k][0] -= size[1] + size[0] * 1.5
    for k in pos2:
        pos2[k] *= size[2]
        pos2[k][0] += size[2] + size[0] * 1.5
    for k in pos3:
        pos3[k] *= size[3]
        pos3[k][1] += size[3] + size[0] * 2
    for k in pos4:
        pos4[k] *= size[4]
        pos4[k][1] -= size[4] + size[0] * 2

    pos = pos0
    pos.update(pos1)
    pos.update(pos2)
    pos.update(pos3)
    pos.update(pos4)
    plt.figure(figsize=(8, 8))
    plt.axis("equal")
    plt.title(info)
    nx.draw_networkx(g,
                     pos=pos,
                     with_labels=False,
                     node_size=25,
                     edge_color="#D8D8D8")
    plt.show()


def plot_in_degree_histogram(G: nx.DiGraph):
    in_degree = np.array([i[1] for i in G.in_degree()])
    plt.title("in degree of each node")
    # plt.xlabel("节点编号(创建顺序)")
    plt.xlabel("node number(order)")
    plt.ylabel("degree")
    plt.plot(in_degree)
    plt.show()

    histogram = np.array([(in_degree == i).sum()
                          for i in range(max(in_degree) + 1)])
    plt.title("in degree histogram")
    plt.xlabel("degree")
    plt.ylabel("number")
    plt.plot(histogram)
    plt.show()

    plt.xlabel("degree(log)")
    plt.ylabel("number(log)")
    plt.xscale("log")
    plt.yscale("log")
    plt.plot(histogram)
    plt.show()


if __name__ == "__main__":
    g = web_page_with_random_edge(1000, 7, 5, 0.7)
    relink(g, 0.01, 2)
    plot_web_structure(g)

    g = web_page_with_random_edge(10000, 10, 5, 0.7)
    relink(g, 0.02, 5)
    plot_in_degree_histogram(g)
